# -*- coding: utf-8 -*-
# @Author: long
# @Date:   2017-12-18 06:20:59
# @Last Modified by:   ystlong
# @Last Modified time: 2017-12-23 08:50:13

import sys
import signal
import uuid

import socket
# from ServerHandle import *
from HookUtil import ProxyThreadPool, delay
import HookProxy
import HookUtil

# target: 被代理的目标， 被访问的目的端
# client: 代理局域网服务， 桥接target与server
# server: 代理公网服务， 桥接client与source
# source: 访问发起源

def new_listen_socket(addr, count, timeout):
	sock = None
	for i in xrange(timeout):
		try:
			sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			sock.bind(addr)
			sock.listen(3)
			break
		except Exception as e:
			HookUtil.clean_socket(sock)
			sock = None
			print(repr(e), "try again: ", i)
			HookUtil.delay(1000)
	return sock

class HookServer(object):
	"""docstring for HookServer"""
	def __init__(self, host, port):
		super(HookServer, self).__init__()
		# self.host = ''
		# self.port = 12345
		self.addr = (host, port)
		self.buffsize = 2048
		self.on = True
		self.connections = {}
		self.connect = False

	def start(self):
		# retry 5 minutes
		hookserver = new_listen_socket(self.addr, 3, 5*60)
		self.hookserver = hookserver
		if hookserver == None:
			print "cann't bind socket ", self.addr
			return
		# hookserver = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		# hookserver.bind(self.addr)
		# hookserver.listen(3)
		while self.on:
			print("HookServer: wait for connection ...")
			client_socket, client_addr = hookserver.accept()
			self.dispatch_socket(client_socket)
		hookserver.close()

	def dispatch_socket(self, socket):
		data = socket.recv(self.buffsize)
		if data.strip().startswith("hook-socket:"):
			key = data.split(":")[1].strip()
			self.connections[key] = socket
			socket.send("hook-ack")
		elif data.strip() == "hook-master":
			self.connect = True
			socket.send("hook-ack")
			self.create_connect(socket)
		elif data == "hook-heart":
			socket.send("hook-heart")
		else:
			# invaild socket type
			socket.send("hook-unknown")
			HookUtil.clean_socket(socket)


	def create_connect(self, socket):
		def runable():
			while self.connect and self.on:
				keys = self.connections.keys()
				for key in keys:
					try:
						if self.connections[key] == None:
							socket.send("hook-socket: " + key)
							# maybe new socket connect
							data = socket.recv(self.buffsize)
							print "=====create_connect===", key, data
							if data == "hook-ack":
								pass
					except KeyError as k:
						pass
					except Exception as e:
						raise e
				if len(self.connections) <= 0:
					delay(1000)
			self.connect = False
		ProxyThreadPool.new_thread(runable)

	def get_client_socket(self, key):
		if self.connections.has_key(key):
				conn = self.connections[key]
				HookUtil.clean_socket(conn)
		self.connections[key] = None
		timeout = 0
		# 超时20s后还未连接成功断开
		while self.connections[key] == None and timeout < 20:
			print "wait socket connect, key: ", key, timeout
			timeout += 1
			delay(100)
		socket = self.connections[key]
		del self.connections[key]
		return socket

	def wheel_source_request(self, source_socket, source_addr):
		def source_request_thread_handle():
			print("gate source request addr", source_addr)
			key = str(source_addr)
			client_socket = self.get_client_socket(key)
			# 没有建立链接
			if client_socket == None:
				print "no client connect"
				HookUtil.clean_socket(source_socket)
				return
			hproxy = HookProxy.HookProxy(source_socket, client_socket)
			hproxy.start()
		ProxyThreadPool.new_thread(source_request_thread_handle)

	def close(self):
		HookUtil.clean_socket(self.hookserver)

class GateServer(object):
	"""网关服务"""
	def __init__(self, hookserver, host, port):
		super(GateServer, self).__init__()
		# self.host = '127.0.0.1'
		# self.port = 1888
		self.addr = (host, port)
		self.hookserver = hookserver
		self.on = True

	def start(self):
		# retry 5 minutes
		self.gate_listen_socket = new_listen_socket(self.addr, 3, 5*60)
		if self.gate_listen_socket == None:
			print "cann't bind socket ", self.addr
			return
		while self.on:
			print("self.gate_listen_socket: wait for connection ...")
			source_socket, source_addr = self.gate_listen_socket.accept()
			try:
				self.hookserver.wheel_source_request(source_socket, source_addr)
			except Exception as e:
				print repr(e)
				raise
		HookUtil.clean_socket(self.gate_listen_socket)

	def close(self):
		self.on = False
		HookUtil.clean_socket(self.gate_listen_socket)
		# self.self.gate_listen_socket.close()

def main():
	if len(sys.argv) != 3:
		print "usage: %s listen_host:port hook_host:port"%(sys.argv[0])
		sys.exit(-1)
	gate_host, gate_port = sys.argv[1].split(":")
	hook_host, hook_port = sys.argv[2].split(":")

	gate_port = int(gate_port)
	hook_port = int(hook_port)

	hs = HookServer(hook_host, hook_port)
	gs = GateServer(hs, gate_host, gate_port)
	ProxyThreadPool.new_thread(hs.start)
	ProxyThreadPool.new_thread(gs.start)

	def exit_handle(signum, frame):
		hs.close()
		gs.close()
		ProxyThreadPool.close()
	signal.signal(signal.SIGINT, exit_handle)
	ProxyThreadPool.wait()


def test():
	sys.argv.append("localhost:1888")
	sys.argv.append("localhost:1887")


if __name__ == '__main__':
	# test()
	main()
